#! /Pub/Users/wangyk/software/pixi/env/omicverse/.pixi/envs/ov/bin/python
ATTRIBUTE_NAME_CELL_IDENTIFIER = "CellID"
ATTRIBUTE_NAME_GENE = "Gene"

def save_df_as_loom(df, fname: str) -> None:
    """
    Save a pandas DataFrame as a single-layer loom file.

    This function takes a pandas DataFrame, where rows represent cells and columns
    represent genes, and saves it as a loom file. Loom files are efficient for
    storing large single-cell datasets, especially sparse matrices. The function
    transposes the DataFrame before saving to ensure that in the loom file,
    rows represent genes and columns represent cells, which is the standard
    orientation for loom files in single-cell genomics.

    The function automatically adds row attributes (gene names) and column
    attributes (cell identifiers) to the loom file, using the DataFrame's
    index as cell identifiers and columns as gene names.

    Args:
        df: A 2-dimensional pandas DataFrame representing the expression matrix.
            Rows should correspond to cells, and columns should correspond to genes.
        fname: The desired filename (including path) for the output loom file.
               The file extension '.loom' will be automatically appended if not present.

    Returns:
        None. This function saves a loom file to disk and does not return any value.

    Raises:
        AssertionError: If the input DataFrame `df` is not 2-dimensional.

    Example:
        假设你有一个 pandas DataFrame 叫做 `expression_df`，你想把它保存为 'my_loom_file.loom'
        >>> save_df_as_loom(expression_df, 'my_loom_file.loom')
    """
    import loompy as lp
    
    assert df.ndim == 2, "Input DataFrame must be 2-dimensional (cells x genes)."
    # The orientation of the loom file is always:
    #   - Columns represent cells or aggregates of cells
    #   - Rows represent genes
    column_attrs = {
        ATTRIBUTE_NAME_CELL_IDENTIFIER: df.index.values.astype("str"),
    }
    row_attrs = {
        ATTRIBUTE_NAME_GENE: df.columns.values.astype("str"),
    }
    # Ensure the filename ends with '.loom'
    if not fname.endswith(".loom"):
        fname += ".loom"

    lp.create(
        filename=fname, layers=df.T.values, row_attrs=row_attrs, col_attrs=column_attrs
    )


def save_raw_x_2_loom(adata_path: str, output_loom_file: str) -> None:
    """
    Convert the raw count matrix (adata.raw.X) of an AnnData object to a loom file.

    This function reads an AnnData object from an .h5ad file, extracts the raw
    count matrix from `adata.raw.X`, converts it to a pandas DataFrame, and
    then saves it as a loom file using the `save_df_as_loom` function.

    The cell identifiers (column attributes in the loom file) are taken from
    `adata.obs.index`, and gene names (row attributes) are taken from
    `adata.raw.var.index`.

    Args:
        adata_path: Path to the input AnnData file (.h5ad).
        output_loom_file: Path to the output loom file (.loom).

    Returns:
        None. This function saves a loom file to disk and does not return any value.

    Example:
        假设你有一个 AnnData 文件 'my_adata.h5ad'，你想把它的 raw count matrix
        保存为 'raw_counts.loom'
        >>> save_raw_x_2_loom('my_adata.h5ad', 'raw_counts.loom')
    """
    
    import pandas as pd
    import scanpy
    import numpy as np
    
    od = output_loom_file
    adata = scanpy.read_h5ad(adata_path)
    
    if hasattr(adata.raw, 'X'):
        df_ = pd.DataFrame(np.array(adata.raw.X)) # Convert sparse matrix to dense for DataFrame
    else:
        df_ = pd.DataFrame(np.array(adata.X)) # Convert sparse matrix to dense for DataFrame
        
    df_.index = adata.obs.index
    
    if hasattr(adata.raw, 'X'):
        df_.columns = adata.raw.var.index
    else:
        df_.columns = adata.var.index

    save_df_as_loom(df=df_, fname=od)


if __name__ == "__main__":
    import argparse
    wrp = '''
    Convert AnnData's raw count matrix (adata.raw.X) to a loom file.

    This script reads an AnnData object from an .h5ad file, extracts the raw
    count matrix, and saves it as a loom file. The output loom file is compatible
    with various single-cell analysis tools, including pySCENIC.

    The cell IDs and gene names are extracted from the AnnData object's 
    `adata.obs.index` and `adata.raw.var.index` respectively.

    Example Usage:
        save_raw_x_2_loom.py -a adata_input.h5ad -o output.loom
    '''
    parser = argparse.ArgumentParser(prog='save_raw_x_2_loom',
                                     description=wrp,
                                     formatter_class=argparse.RawTextHelpFormatter)
    parser.add_argument('-a', '--adata_path',
                        help="Path to the input AnnData file (.h5ad)",
                        required=True,
                        metavar='')
    parser.add_argument('-o', '--output_loom_file',
                        help="Path to the output loom file (.loom)",
                        required=True,
                        metavar='')


    args = parser.parse_args()
    save_raw_x_2_loom(adata_path=args.adata_path, output_loom_file=args.output_loom_file)